import cv2
import numpy as np

# 左上角和右下角去除水印
def remove_watermarks(input_path, output_path):
    # 读取图片
    img = cv2.imread(input_path)
    height, width = img.shape[:2]

    # 创建全黑掩码（单通道）
    mask = np.zeros((height, width), dtype=np.uint8)

    # 标记左上角水印区域（需根据实际水印大小调整坐标）
    top_left_x_start, top_left_y_start = 0, 0  # 左上角起点
    top_left_x_end, top_left_y_end = 80, 40  # 假设水印宽100像素，高50像素
    mask[top_left_y_start:top_left_y_end, top_left_x_start:top_left_x_end] = 255

    # 标记右下角水印区域
    bottom_right_x_start, bottom_right_y_start = width - 80, height - 40  # 右下角起点
    bottom_right_x_end, bottom_right_y_end = width, height
    mask[bottom_right_y_start:bottom_right_y_end, bottom_right_x_start:bottom_right_x_end] = 255

    # 使用图像修复算法
    result = cv2.inpaint(img, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
    # 保存结果
    cv2.imwrite(output_path, result)
    print(f"左下角-右下角水印已去除，结果保存至: {output_path}")


# 右下角去除水印
def remove_watermark(input_path, output_path):
    # 读取图片
    img = cv2.imread(input_path)
    height, width = img.shape[:2]
    # 手动定义水印区域（右下角，需根据实际调整坐标）
    x_start, y_start = width - 100, height - 50  # 水印右下角坐标
    x_end, y_end = width, height  # 水印区域结束坐标
    # 创建掩码（标记水印区域为白色，其余为黑色）
    mask = np.zeros((height, width), np.uint8)
    mask[y_start:y_end, x_start:x_end] = 255  # 白色区域代表需修复部分
    # 使用图像修复算法
    result = cv2.inpaint(img, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
    # 保存结果
    cv2.imwrite(output_path, result)
    print(f"右下角水印已去除，结果保存至: {output_path}")

# 读取图片二进制数据去水印
def remove_watermarks_in_memory(image_bytes):
    """
    直接在内存中去水印
    :param image_bytes: 图片的二进制数据(bytes)
    :return: 处理后的图片二进制数据(bytes)
    """
    # 将二进制图片数据转为numpy数组
    img_array = np.frombuffer(image_bytes, np.uint8)

    # 解码为OpenCV图像格式
    img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError("无法解码图片数据")

    height, width = img.shape[:2]

    # 创建全黑掩码（单通道）
    mask = np.zeros((height, width), dtype=np.uint8)

    # 标记左上角水印区域（示例值，需根据实际情况调整）
    top_left_x_start, top_left_y_start = 0, 0
    top_left_x_end, top_left_y_end = 90, 40
    mask[top_left_y_start:top_left_y_end, top_left_x_start:top_left_x_end] = 255

    # 标记右下角水印区域
    bottom_right_x_start, bottom_right_y_start = width - 150, height - 50
    bottom_right_x_end, bottom_right_y_end = width, height
    mask[bottom_right_y_start:bottom_right_y_end, bottom_right_x_start:bottom_right_x_end] = 255

    # 使用图像修复算法
    result = cv2.inpaint(img, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)

    # 将处理后的图像编码为JPEG格式的二进制数据
    success, encoded_img = cv2.imencode('.jpg', result)
    if not success:
        raise RuntimeError("图片编码失败")

    return encoded_img.tobytes()


# remove_watermarks("img/" + "input.jpg", "output.jpg")